# Copyright (C) 2009,2010,2011,2012  Marco Restelli
#
# This file is part of:
#   FEMilaro -- Finite Element Method toolkit
#
# FEMilaro is free software; you can redistribute it and/or modify it
# under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# FEMilaro is distributed in the hope that it will be useful, but
# WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
# General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with FEMilaro; If not, see <http://www.gnu.org/licenses/>.
#
# author: Marco Restelli                   <marco.restelli@gmail.com>

"""
Compute the "exact" solution of the linear inertia-gravity wave
problem for the nonhydrostatic case.

Notice that the horizontal direction is treated by Fourier series:
this is exact if the domain is periodic, while it is an approximation
if the domain is infinite - in this latter case one should take the
domain as large as possible.

Notice that minor modifications would be required to handle the
Boussinesq and the hydrostatic cases.
"""

import time
import mpmath as mpm

# Set the working precision
mpm.mp.dps = 25 # decimal digits

# Define the problem constants
N = mpm.mpf('0.01')   # Brunt-Vaisalla freq.
T = mpm.mpf('224.0')  # temperature
g = mpm.mpf('9.80616')
cp = mpm.mpf('1004.67')
R  = mpm.mpf('287.17')
cv = cp-R
gamma = cp/cv
kappa = R/cp
p_s = mpm.mpf('1.0e5') # sea level reference pressure
a = mpm.sqrt(gamma*R*T)

Lx = mpm.mpf('300e3') # domain horizontal size
Lz = mpm.mpf('10e3')  # domain vertical size
xc = mpm.mpf('100e3')    # perturbation center
aa = mpm.mpf('5e3')      # perturbation width
dtheta = mpm.mpf('0.01') # perturbation amplitude
U = mpm.mpf('20.0')      # mean flow velocity

t = mpm.mpf('3000.0') # time
z = mpm.mpf('4.5e3')  # height
Nx = 12000             # number of points in x direction
Nk = 6000             # number of Fourier modes

r1 = mpm.mpf(1)
r2 = mpm.mpf(2)
r4 = mpm.mpf(4)
l  = mpm.pi/Lz
I  = [-Lx/r2,Lx/r2]  # x interval

# Define the sampling points
x = [ Lx*(-r1/r2+mpm.mpf(i)/mpm.mpf(Nx-1)) for i in range(Nx) ]

# Simplified version of mpm.fourier (see
# mpmath/calculus/approximation.py): Fourier coefficients are
# evaluated with 2N equally spaced nodes.
def Fourier(N,f):
  k1 = r2*mpm.pi
  x  = [-r1/r2+mpm.mpf(i)/mpm.mpf(2*N) for i in xrange(2*N)]
  fx = [f(Lx*xi) for xi in x]
  F = []
  for ki in xrange(N):
    kx = k1*mpm.mpf(ki)
    F.append( mpm.fsum( fi*mpm.cos(kx*xi) for xi,fi in zip(x,fx) ) )
    F[-1] /= mpm.mpf(N)
  F[0] /= r2 # scale the first coefficient (mean value)
  return F

# Simplified version of mpm.fourierval (see
# mpmath/calculus/approximation.py): Fourier series is evaluated on
# arbitrary points following the definition.
def invFourier(x,F):
  k1 = r2*mpm.pi
  f = []
  for xi in x:
    kx = k1*(xi/Lx)
    f.append( mpm.fsum( F[ki]*mpm.cos(kx*mpm.mpf(ki))   \
                          for ki in xrange(len(F))    ) )
  return f

T0 = time.clock()

# Fourier coefficients of the initial condition (these could be
# obtained analytically, but then one should pay attention to various
# normalization constants; moreover, a small error would be introduced
# because the domain is limited and periodic, while the analytic
# coefficients are computed for an infinite domain). Notice as well
# that these syntax only captures the cosine coefficients, which is
# correct for the standard initial condition.
def agnesi(x):
  return dtheta/((x/aa)**2+r1)*mpm.sin(l*z)
THk0 = Fourier( Nk , agnesi )

T1 = time.clock(); print('Fourier time: %f'%(T1-T0))
  
# Preliminary check: compare the initial datum with its Fourier
# series.
TH0 = invFourier(x,THk0)

print('Inverse Fourier time: %f'%(time.clock()-T1))

err0 = [ mpm.fabs(agnesi(x[i])-TH0[i]) for i in xrange(len(x)) ]

T1 = time.clock(); print('Total setup time: %f'%(T1-T0))

print('Maximum initial error: %e' % max( float(ei) for ei in err0 ))
import pylab as py
py.figure(1)
py.plot([float(xi) for xi in x],[float(ei) for ei in err0],'go-')
py.title("Initial error")
py.figure(2)
py.title("Initial datum and Fourier series")
py.plot(                     \
  [float(xi) for xi in x] ,  \
  [float(ti) for ti in TH0], \
  'mo-', label='Fourier series')
py.plot(                                          \
  [float(xi) for xi in x] ,                       \
  [float(agnesi(x[i])) for i in xrange(len(x))] , \
  'k-', label='initial datum')
py.legend(loc='upper right')

# Now the real computations
T0 = time.clock()
THk = []
N2 , a2 = N**2 , a**2
k1 = r2*mpm.pi/Lx
for ki in range(Nk):
  k = k1*mpm.mpf(ki)
  k2 = k**2
  PHI = a2*(k2+l**2) + N2
  tmp = mpm.sqrt(PHI**2 - r4*a2*k2*N2)
  PSIs2 = PHI + tmp
  PSIg2 = PHI - tmp
  den = r2*tmp
  Os = mpm.sqrt(PSIs2/r2)
  Og = mpm.sqrt(PSIg2/r2)
  theta = THk0[ki]                                 \
   * ( (mpm.mpf(2)*N2-PSIg2)/den * mpm.cos(Os*t)   \
      +(PSIs2-mpm.mpf(2)*N2)/den * mpm.cos(Og*t) )
  THk.append(theta)

TH = invFourier((xi-U*t+(Lx/r2-xc) for xi in x),THk)

T1 = time.clock()
print('Computation time: %f'%(T1-T0))

py.figure(3)
py.plot(                            \
  [float(xi+xc) for xi in x] ,  \
  [float(ti) for ti in TH0],        \
  'k-', label='initial condition')
py.plot(                     \
  [float(xi+Lx/r2) for xi in x] ,  \
  [float(ti) for ti in TH], \
  'r-', label='time %4.0fs'%t)
py.legend(loc='upper right')
py.show()

